from pyflix.datasets import RatedDataset

CLIP_RES = .025

tr = RatedDataset('/netflix/database/training_set')
pr = RatedDataset('/netflix/database/probe_set')

class Machine(object):
  name = "machine"
  clip_low = 1.0
  clip_high = 5.0
  filters = None
  
  def add_filter(self, filter):
    if self.filters is None:
      self.filters = [filter]
    else:
      self.filters.append(filter)
  
  def find_clip(self, outset=pr, res=CLIP_RES, quiet=False):
    if not quiet:
      print "finding low clip"
    clips = [1. + i * res for i in range(int(1.05/res))]
    se_s = [0.] * len(clips)
    rs = 0
    for mid, uid, r in outset.iterMovieIDsUserIDsRatings():
      rs += 1
      est = self.get_rating(mid, uid)
      for i, clip in enumerate(clips):
        se_s[i] += (max(clip, est) - r) ** 2
    if not quiet:
      for clip, se in zip(clips, se_s):
        print "  %s: rmse = %s" % (clip, (se / rs) ** .5)
    self.clip_low = clips[se_s.index(min(se_s))]
    
    if not quiet:
      print "finding high clip"
    clips = [5. - i * res for i in range(int(1.05/res))]
    se_s = [0.] * len(clips)
    rs = 0
    for mid, uid, r in outset.iterMovieIDsUserIDsRatings():
      rs += 1
      est = self.get_rating(mid, uid)
      for i, clip in enumerate(clips):
        se_s[i] += (min(clip, est) - r) ** 2
    if not quiet:
      for clip, se in zip(clips, se_s):
        print "  %s: rmse = %s" % (clip, (se / rs) ** .5)
    self.clip_high = clips[se_s.index(min(se_s))]
  
  def get_rating(self, mid, uid):
    r = self.est_rating(mid, uid)
    r = max(min(self.clip_high, r), self.clip_low)
    if self.filters:
      for filter in self.filters:
        r = filter.rectify(r, mid, uid)
    return r
  
  
  def calc_rmse(self, insamp=False, quiet=False, inset=tr, outset=pr):
    if insamp:
      rs = 0
      err_sq = 0.
      for mid, uid, r in inset.iterMovieIDsUserIDsRatings():
        rs += 1
        err_sq += (self.est_rating(mid, uid)-r)**2
      rmse = ((err_sq / rs) ** .5)
      if not quiet:
        print "insample RMSE = %s" % rmse,
        print "(%s %%)" % (100 * (1 - rmse/.9525))
    rs = 0
    err_sq = 0.
    for mid, uid, r in outset.iterMovieIDsUserIDsRatings():
      rs += 1
      err_sq += (self.get_rating(mid, uid)-r)**2
    rmse = ((err_sq / rs) ** .5)
    if not quiet:
      print "outsample RMSE = %s" % rmse,
      print "(%s %%)" % (100 * (1 - rmse/.9525))
    return rmse
  
      